package agent

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"strings"

	cfg "openmanus/config"
	"openmanus/llm"
	"openmanus/llm/history"
	"openmanus/tool"

	"github.com/axcom/ninego/class"
	logger "github.com/axcom/ninego/log"

	mcpclient "github.com/mark3labs/mcp-go/client"
	"github.com/mark3labs/mcp-go/mcp"
)

// 实现了工具调用机制，使智能体能够使用各种工具来完成任务。
type ToolCallAgent struct {
	ReActAgent
	MaxObserve       interface{}
	MCPClients       map[string]*mcpclient.Client // MCP工具
	AvailableTools   *tool.ToolCollection         // 可用工具集合
	SpecialToolNames []string                     //特殊工具名称列表
	ToolChoices      TOOL_CHOICE_TYPE             //工具选择模式（"none"、"auto"、"required"）

	messageContent []history.ContentBlock //当前处理的信息
	ToolCalls      []llm.ToolCall         // 工具调用列表
	ToolResults    []history.ContentBlock //工具调用返回信息
}

const (
	SYSTEM_PROMPT    = "You are an agent that can execute tool calls"
	NEXT_STEP_PROMPT = "If you want to stop interaction, use `terminate` tool/function call."
)

func NewToolCallAgent(name string, llm *LLMClient) *ToolCallAgent {
	agent := &ToolCallAgent{
		ReActAgent: ReActAgent{
			BaseAgent: BaseAgent{
				Name:           "toolcall",
				Description:    "an agent that can execute tool calls.",
				SystemPrompt:   SYSTEM_PROMPT,
				NextStepPrompt: NEXT_STEP_PROMPT,
				MaxSteps:       30,
				Memory:         &history.Memory{Messages: make([]history.HistoryMessage, 0)},
				LLM:            llm,
			},
		},
		MaxObserve:       nil,
		AvailableTools:   &tool.ToolCollection{},
		ToolChoices:      ToolChoiceAuto,
		SpecialToolNames: []string{"Terminate"},
		//ToolCalls:        []llm.ToolCall{},
	}

	return agent
}

func (agent *ToolCallAgent) Initialize_mcp_servers(mcpConfig *cfg.MCPSettings) {
	clients, err := createMCPClients(mcpConfig)
	logger.Println(mcpConfig)
	if err != nil {
		panic(fmt.Errorf("error creating MCP clients: %v", err))
	}
	agent.MCPClients = clients

	for serverName, mcpClient := range agent.MCPClients {
		logger.Info("Server connected", "name", serverName)

		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		toolsResult, err := mcpClient.ListTools(ctx, mcp.ListToolsRequest{})
		cancel()
		if err != nil {
			logger.Error(
				"Error fetching tools",
				"server",
				serverName,
				"error",
				err,
			)
			continue
		}
		if len(toolsResult.Tools) == 0 {
			fmt.Println(" ", "No tools available")
		} else {
			for _, t := range toolsResult.Tools {
				// Add the tool with its description as a nested list
				namespacedName := fmt.Sprintf("%s__%s", serverName, t.Name)
				fmt.Println(" ", t.Description)
				serverTool := tool.NewMCPClientTool(
					namespacedName,
					t.Description,
					llm.Schema{
						Type:       t.InputSchema.Type,
						Properties: t.InputSchema.Properties,
						Required:   t.InputSchema.Required,
					},
					mcpClient,
				)
				// 线程安全地更新工具映射
				agent.AvailableTools.AddTool(serverTool)
			}
		}
	}
}

func (agent *ToolCallAgent) Disconnect_mcp_server() {
	// 清理MCP资源
	logger.Info("Shutting down MCP servers...")
	for name, client := range agent.MCPClients {
		if err := client.Close(); err != nil {
			logger.Error("Failed to close server", "name", name, "error", err)
		} else {
			logger.Info("Server closed", "name", name)
		}
	}
}

// 使用LLM决定使用哪些工具
func (agent *ToolCallAgent) Think() (bool, error) {
	ret, handled := agent.Inherited()
	if handled {
		return class.SC[bool](ret[0]), class.SC[error](ret[1])
	}

	fmt.Println("Think...?")

	if agent.NextStepPrompt != "" {
		agent.UpdateMemory(USER, agent.NextStepPrompt)
		agent.NextStepPrompt = "" // 清空它，只添加1次（在后续可改变它的内容）
	}

	llmMessages := make([]llm.Message, len(agent.Memory.Messages))
	for i := range agent.Memory.Messages {
		llmMessages[i] = &(agent.Memory.Messages)[i]
	}

	// 获取带工具选项的响应
	var response llm.Message
	var err error
	response, err = agent.LLM.AskTool(llmMessages, "", agent.AvailableTools.ToParams())
	if err != nil {
		//fmt.Println(err)
		/*/ 检查是否是TokenLimitExceeded错误
		if _, ok := err.(*TokenLimitExceeded); ok {
			logger.Error(fmt.Sprintf("🚨 Token limit error: %v", err))
			a.Memory.AddMessage(
				schema.NewAssistantMessage(
					fmt.Sprintf("Maximum token limit reached, cannot continue execution: %v", err),
					"",
				),
			)
			a.State = schema.AgentStateFINISHED
			return false, nil
		}*/
		return false, err
	}
	agent.ToolCalls = response.GetToolCalls()
	logger.Printf("✨ %s's thoughts: %s\n", agent.Name, response.GetContent())
	logger.Printf("🛠️ %s selected %d tools to use\n", agent.Name, len(response.GetToolCalls()))
	if len(response.GetToolCalls()) > 0 {
		toolNames := make([]string, len(response.GetToolCalls()))
		for i, call := range response.GetToolCalls() {
			toolNames[i] = call.GetName()
		}
		logger.Printf("🧰 Tools being prepared: %s\n", strings.Join(toolNames, ", "))
		logger.Printf("🔧 Tool arguments: %v\n", response.GetToolCalls()[0])
	}

	if agent.ToolChoices == ToolChoiceNone {
		if len(response.GetToolCalls()) > 0 {
			logger.Printf("🤔 Hmm, %s tried to use tools when they weren't available!\n", agent.Name)
		}
		return false, nil
	}

	// Log usage statistics if available
	inputTokens, outputTokens := response.GetUsage()
	if inputTokens > 0 || outputTokens > 0 {
		logger.Info("Usage statistics",
			"input_tokens", inputTokens,
			"output_tokens", outputTokens,
			"total_tokens", inputTokens+outputTokens)
	}

	// 获取内容
	content := response.GetContent()
	// Add text content
	if content != "" {
		agent.messageContent = append(agent.messageContent, history.ContentBlock{
			Type: "text",
			Text: content,
		})
	}
	logger.Debug("ask-GPT-role=", response.GetRole(), response.GetContent(), len(agent.ToolCalls))

	// 处理不同的tool_choices模式
	switch agent.ToolChoices {
	case ToolChoiceNone:
		if len(agent.ToolCalls) > 0 {
			logger.Warn(fmt.Sprintf("🤔 Hmm, %s tried to use tools when they weren't available!", agent.Name))
		}
		if content != "" {
			//a.Memory.AddMessage(schema.NewAssistantMessage(content, ""))
			agent.messageContent = append(agent.messageContent, history.ContentBlock{
				Type: "text",
				Text: response.GetContent(),
			})
			return true, nil
		}
		return false, nil
	case ToolChoiceRequired:
		// 创建助手消息
		var assistantMsg history.ContentBlock
		if len(agent.ToolCalls) > 0 {
			//exec-tool--> assistantMsg = schema.NewAssistantMessageWithToolCalls(content, a.ToolCalls)
		} else {
			//assistantMsg = schema.NewAssistantMessage(content, "")
			assistantMsg = history.ContentBlock{
				Type: "text",
				Text: response.GetContent(),
			}
			agent.messageContent = append(agent.messageContent, assistantMsg) //a.Memory.AddMessage(assistantMsg)
		}
		return true, nil
	case ToolChoiceAuto:
		// 创建助手消息
		var assistantMsg history.ContentBlock
		if len(agent.ToolCalls) > 0 {
			//exec-tool--> assistantMsg = schema.NewAssistantMessageWithToolCalls(content, a.ToolCalls)
		} else {
			//assistantMsg = schema.NewAssistantMessage(content, "")
			assistantMsg = history.ContentBlock{
				Type: "text",
				Text: response.GetContent(),
			}
			agent.messageContent = append(agent.messageContent, assistantMsg) //a.Memory.AddMessage(assistantMsg)
		}

		// 如果没有工具调用但有内容
		if len(agent.ToolCalls) == 0 {
			//return content != "", nil
			if len(agent.Memory.Messages) > 1 { //上上次工具调用后，本次无工具调用 --> 结束
				if agent.Memory.Messages[len(agent.Memory.Messages)-1-1].Role == TOOL {
					return false, nil
				}
			}
		}
		//return true, nil
	}
	return len(agent.ToolCalls) > 0, nil

}

// 执行工具调用并处理结果
func (agent *ToolCallAgent) Act() (string, error) {
	ret, handled := agent.Inherited()
	if handled {
		return class.SC[string](ret[0]), class.SC[error](ret[1])
	}

	if len(agent.ToolCalls) == 0 {
		if agent.ToolChoices == ToolChoiceRequired {
			return "", errors.New("Tool calls required but none provided")
		}
		if len(agent.Memory.Messages) > 0 {
			// 返回最后一条消息的内容
			return agent.Memory.Messages[len(agent.Memory.Messages)-1].GetContent(), nil
		}
		return "No content or commands to execute", nil
	}

	// 清空本次 messageContent / ToolResults
	agent.messageContent = []history.ContentBlock{}
	agent.ToolResults = []history.ContentBlock{}

	results := []string{}
	for _, command := range agent.ToolCalls {
		result, err := agent.ExecuteTool(command)
		if err != nil {
			return "", err
		}

		// 限制观察结果的长度
		/*if agent.MaxObserve != nil {
			if maxObserve, ok := agent.MaxObserve.(int); ok {
				result = result[:maxObserve]
			}
		}*/

		logger.Printf("🎯 Tool '%s' completed its mission! Result: %s\n", command.GetName(), result)
		results = append(results, result)
	}

	// 添加工具调用到agent.Memory
	agent.Memory.AddMessage(history.HistoryMessage{
		Role:    ASSISTANT, //message.GetRole(),
		Content: agent.messageContent,
	})

	// 添加工具响应到agent.Memory
	agent.Memory.AddMessage(history.HistoryMessage{
		Role:    TOOL,
		Content: agent.ToolResults,
	})

	// 将结构体转换为JSON
	jsonData, _ := json.Marshal(agent.Memory)
	logger.Debug("<Memory=>", string(jsonData), "\n<=Memory>")

	return strings.Join(results, "\n\n"), nil
}

// 执行单个工具调用
func (agent *ToolCallAgent) ExecuteTool(toolCall llm.ToolCall) (string, error) {
	logger.Println("🔧 Using tool", "name", toolCall.GetName())

	input, err := json.Marshal(toolCall.GetArguments())
	if err != nil {
		fmt.Printf("Error parsing tool arguments: %v\n", err)
		return "", err
	}
	agent.messageContent = append(agent.messageContent, history.ContentBlock{
		Type:  "tool_use",
		ID:    toolCall.GetID(),
		Name:  toolCall.GetName(),
		Input: input,
	})

	// 获取工具
	_, exists := agent.AvailableTools.GetTool(toolCall.GetName())
	if !exists {
		return "", fmt.Errorf("工具 '%s' 不存在", toolCall.GetName())
	}

	logger.Printf("执行工具: %s, 参数: %v\n", toolCall.GetName(), toolCall.GetArguments())

	// 执行工具
	result, err := agent.AvailableTools.Execute(toolCall.GetName(), toolCall.GetArguments())
	if err != nil {
		// 记录错误结果
		errorMsg := fmt.Sprintf("工具 '%s' 执行失败: %v", toolCall.GetName(), err)
		agent.ToolResults = append(agent.ToolResults, history.ContentBlock{
			Type:      "tool_result",
			ToolUseID: toolCall.GetID(),
			Content: []history.ContentBlock{{
				Type: "text",
				Text: errorMsg,
			}},
		})
		return errorMsg, err
	}

	// 记录成功结果
	if result.Output != nil {
		agent.ToolResults = append(agent.ToolResults, history.ContentBlock{
			Type:      "tool_result",
			ToolUseID: toolCall.GetID(),
			Content:   result.Output,
			Text:      result.Text,
		})
	} else {
		agent.ToolResults = append(agent.ToolResults, history.ContentBlock{
			Type:      "tool_result",
			ToolUseID: toolCall.GetID(),
			Text:      result.Text,
		})
	}

	// 处理特殊工具
	if result, err := agent.HandleSpecialTool(toolCall); err != nil {
		return fmt.Sprintf("Error handling special tool: %s %v", result, err), nil
	}
	if agent._shouldFinishExecution() {
		// 设置代理状态为已完成
		logger.Info(fmt.Sprintf("🏁 Special tool '%s' has completed the task!", toolCall.GetName()))
		agent.State = FINISHED
	}

	//return result.Content, nil
	return result.StringResult(), nil
}

// _isSpecialTool 检查是否是特殊工具
func (agent *ToolCallAgent) _isSpecialTool(toolName string) bool {
	// 特殊工具名称列表
	for _, specialTool := range agent.SpecialToolNames {
		if toolName == specialTool {
			return true
		}
	}
	return false
}

// _handleSpecialTool 处理特殊工具
func (a *ToolCallAgent) HandleSpecialTool(toolCall llm.ToolCall) (string, error) {
	if !a._isSpecialTool(toolCall.GetName()) {
		return "", nil
	}
	switch toolCall.GetName() {
	case "terminate":
		// 处理终止工具
		reason := "无特殊原因"
		if reasonParam, ok := toolCall.GetArguments()["reason"].(string); ok {
			reason = reasonParam
		}
		a.State = FINISHED
		return fmt.Sprintf("终止: %s", reason), nil
	default:
		return "", fmt.Errorf("未实现的特殊工具: %s", toolCall.GetName())
	}
}

// _shouldFinishExecution 检查是否应该结束执行
func (a *ToolCallAgent) _shouldFinishExecution() bool {
	// 检查状态
	if a.State == FINISHED || a.State == ERROR {
		return true
	}

	// 检查最后一条消息
	messages := a.Memory.Messages
	if len(messages) == 0 {
		return false
	}

	lastMessage := messages[len(messages)-1]
	// 如果最后一条消息是助手直接回答，应该结束
	return lastMessage.Role == ASSISTANT && lastMessage.GetToolCalls() == nil
}

// Cleanup 清理资源
func (agent *ToolCallAgent) Cleanup() error {
	// 清理工具资源
	if agent.AvailableTools != nil {
		if err := agent.AvailableTools.Cleanup(); err != nil {
			logger.Error("清理工具集合时出错: %v", err)
		}
	}

	// 清理其他资源
	agent.ToolCalls = nil

	ret, ok := agent.Super()()
	if ok {
		return class.SC[error](ret[0])
	} else {
		return nil
	}
}
